#!/usr/bin/env python
import os
import socket
import sys
from datetime import datetime
from pathlib import Path

import numpy as np
import setproctitle
import torch
import wandb

script_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(os.path.dirname(os.path.dirname(script_dir)))
sys.path.append(parent_dir)
from env_creater import MyEnv
from onpolicy.envs.env_wrappers import DummyVecEnv, SubprocVecEnv

"""Train script for MPEs."""
# 设置环境变量
# os.environ["WANDB_API_KEY"] = " "
## 88plantform

from config import config_8, get_config


def parse_args(parser):
    env = "88Env"
    scenario = "simple_spread"
    num_enemys = 1
    num_agents = 3
    algo = "rmappo"
    exp = "check"
    seed_max = 1
    argv = [
        "--env_name",
        f"{env}",
        "--algorithm_name",
        f"{algo}",
        "--experiment_name",
        f"{exp}",
        "--scenario_name",
        f"{scenario}",
        "--num_agents",
        f"{num_agents}",
        "--num_enemys",
        f"{num_enemys}",
        "--seed",
        f"{seed_max}",
        "--n_training_threads",
        "4",
        "--n_rollout_threads",
        "1",
        "--num_mini_batch",
        "10",
        "--episode_length",
        "200",
        "--num_env_steps",
        "10000000",
        "--ppo_epoch",
        "20",
        "--gain",
        "0.15",
        "--lr",
        "7e-4",
        "--critic_lr",
        "5e-4",
        # "--clip_param",
        # "0.1",
        # "--entropy_coef",
        # "0.01",
        # "--gae_lambda",
        # "0.99",
        "--wandb_name",
        "test_wandb_mpe",
    ]
    parser.add_argument(
        "--scenario_name",
        type=str,
        default="3_1Trapping",
        help="Which scenario to run on",
    )

    all_args = parser.parse_known_args(argv)[0]

    return all_args


def is_port_free(port):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.bind(("", port))
            return True
        except OSError:
            return False


def find_free_ports_with_neighbors(num_ports):
    free_ports = []
    port = 1001  # 起始端口，避免使用系统保留端口
    while len(free_ports) < num_ports:
        if all(is_port_free(6100 + port * 10 + i) for i in range(4)):
            free_ports.append(port)
        port += 1  # 尝试下一个端口
    return free_ports


##
def make_train_env(all_args):
    # 查找100个空闲端口，每个端口和后边紧挨的3个端口都是空闲的
    free_ports = find_free_ports_with_neighbors(all_args.n_rollout_threads)

    def get_env_fn(rank, port):
        def init_env():
            env = MyEnv(all_args, config_8, port)
            env.seed(all_args.seed + rank * 1000)
            return env

        return init_env

    if all_args.n_rollout_threads == 1:
        return DummyVecEnv([get_env_fn(0, 1001)])
    else:
        return SubprocVecEnv(
            [get_env_fn(i, free_ports[i]) for i in range(all_args.n_rollout_threads)]
        )


def make_eval_env(all_args):
    def get_env_fn(rank, port):
        def init_env():
            env = MyEnv(all_args, config_8, port)
            env.seed(all_args.seed * 50000 + rank * 10000)
            return env

        return init_env

    if all_args.n_eval_rollout_threads == 1:
        return DummyVecEnv([get_env_fn(0, 990)])
    else:
        return SubprocVecEnv(
            [get_env_fn(i, 990 + i) for i in range(all_args.n_eval_rollout_threads)]
        )


def main(args):
    all_args = parse_args(get_config())

    if all_args.algorithm_name == "rmappo":
        print("u are choosing to use rmappo, we set use_recurrent_policy to be True")
        all_args.use_recurrent_policy = True
        all_args.use_naive_recurrent_policy = False
    elif all_args.algorithm_name == "mappo":
        print(
            "u are choosing to use mappo, we set use_recurrent_policy & use_naive_recurrent_policy to be False"
        )
        all_args.use_recurrent_policy = False
        all_args.use_naive_recurrent_policy = False
    elif all_args.algorithm_name == "ippo":
        print("u are choosing to use ippo, we set use_centralized_V to be False")
        all_args.use_centralized_V = False
    else:
        raise NotImplementedError

    assert (
        (
            all_args.share_policy == True
            and all_args.scenario_name == "simple_speaker_listener"
        )
        == False
    ), "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."

    # cuda
    if all_args.cuda and torch.cuda.is_available():
        print("choose to use gpu...")
        device = torch.device("cuda:0")
        torch.set_num_threads(all_args.n_training_threads)
        if all_args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        print("choose to use cpu...")
        device = torch.device("cpu")
        torch.set_num_threads(all_args.n_training_threads)

    # run dir
    run_dir = (
        Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[0] + "/results")
        / all_args.env_name
        / all_args.scenario_name
        / all_args.algorithm_name
        / all_args.experiment_name
    )
    if not run_dir.exists():
        os.makedirs(str(run_dir))

    # wandb
    if all_args.use_wandb:
        run = wandb.init(
            config=all_args,
            project=all_args.env_name,
            # entity=all_args.user_name,
            notes=socket.gethostname(),
            name=str(all_args.algorithm_name)
            + "_"
            + datetime.now().strftime("%m_%d_%H_%M_%S"),
            # group=all_args.scenario_name,
            dir=str(run_dir),
            job_type="training",
            reinit=True,
        )
    else:
        if not run_dir.exists():
            curr_run = "run1"
        else:
            exst_run_nums = [
                int(str(folder.name).split("run")[1])
                for folder in run_dir.iterdir()
                if str(folder.name).startswith("run")
            ]
            if len(exst_run_nums) == 0:
                curr_run = "run1"
            else:
                curr_run = "run%i" % (max(exst_run_nums) + 1)
        run_dir = run_dir / curr_run
        if not run_dir.exists():
            os.makedirs(str(run_dir))

    setproctitle.setproctitle(
        str(all_args.algorithm_name)
        + "-"
        + str(all_args.env_name)
        + "-"
        + str(all_args.experiment_name)
        + "@"
        + str(all_args.user_name)
    )

    # seed
    torch.manual_seed(all_args.seed)
    torch.cuda.manual_seed_all(all_args.seed)
    np.random.seed(all_args.seed)

    # env init
    envs = make_train_env(all_args)
    eval_envs = make_eval_env(all_args) if all_args.use_eval else None
    num_agents = all_args.num_agents

    config = {
        "all_args": all_args,
        "envs": envs,
        "eval_envs": eval_envs,
        "num_agents": num_agents,
        "device": device,
        "run_dir": run_dir,
    }

    # run experiments
    if all_args.share_policy:
        from onpolicy.runner.shared.mpe_runner import MPERunner as Runner
    else:
        from onpolicy.runner.separated.mpe_runner import MPERunner as Runner

    runner = Runner(config)
    runner.run()

    # post process
    envs.close()
    if all_args.use_eval and eval_envs is not envs:
        eval_envs.close()

    if all_args.use_wandb:
        run.finish()
    else:
        runner.writter.export_scalars_to_json(str(runner.log_dir + "/summary.json"))
        runner.writter.close()


if __name__ == "__main__":
    # env = "MPE"
    # scenario = "simple_spread"
    # num_enemys = 1
    # num_agents = 3
    # algo = "rmappo"
    # exp = "check"
    # seed_max = 1
    # # 设置 sys.argv 的值
    # argv = [
    #     "train_mpe.py",
    #     "--env_name",
    #     f"{env}",
    #     "--algorithm_name",
    #     f"{algo}",
    #     "--experiment_name",
    #     f"{exp}",
    #     "--scenario_name",
    #     f"{scenario}",
    #     "--num_agents",
    #     f"{num_agents}",
    #     "--num_enemys",
    #     f"{num_enemys}",
    #     "--seed",
    #     f"{seed_max}",
    #     "--n_training_threads",
    #     "1",
    #     "--n_rollout_threads",
    #     "100",
    #     "--num_mini_batch",
    #     "1",
    #     "--episode_length",
    #     "25",
    #     "--num_env_steps",
    #     "20000000",
    #     "--ppo_epoch",
    #     "10",
    #     "--gain",
    #     "0.01",
    #     "--lr",
    #     "7e-4",
    #     "--critic_lr",
    #     "7e-4",
    #     "--wandb_name",
    #     "test_wandb_mpe",
    # ]
    # sys.argv = argv
    main(sys.argv[1:])
